# 提出问题--->文档加载--->文档分割--->embeddings词向量--->存储至向量数据库--->向量数据库根据问题检索数据--->生成含Prompt的问题--->将问题传递给LLM--->LLM给出答案

# 向量数据库与词向量(Vectorstores and Embeddings)

import os
import numpy as np
from langchain.document_loaders import PyPDFLoader
from langchain.embeddings.openai import OpenAIEmbeddings
from langchain.vectorstores import Chroma

api_key = "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"
os.environ['OPENAI_API_KEY'] = "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"


# 在机器学习和自然语言处理（NLP）中，`Embeddings`（嵌入）是一种将类别数据，如单词、句子或者整个文档，转化为实数向量的技术。这些实数向量可以被计算机更好地理解和处理。嵌入背后的主要想法是，相似或相关的对象在嵌入空间中的距离应该很近。
def embeddings_example():
    embedding = OpenAIEmbeddings()
    sentence1_chinese = "我喜欢狗"
    sentence2_chinese = "我喜欢犬科动物"
    sentence3_chinese = "外面的天气很糟糕"
    embedding1 = embedding.embed_query(sentence1_chinese)
    embedding2 = embedding.embed_query(sentence2_chinese)
    embedding3 = embedding.embed_query(sentence3_chinese)
    result1 = np.dot(embedding1, embedding2)
    result2 = np.dot(embedding1, embedding3)
    result3 = np.dot(embedding2, embedding3)

    print(result1)
    print(result2)
    print(result3)


def vectorstores_chroma_example():
    # 1、加载 PDF
    loaders = [
        # 故意添加重复文档，使数据混乱
        PyPDFLoader("docs/cs229_lectures/MachineLearning-Lecture01.pdf"),
        PyPDFLoader("docs/cs229_lectures/MachineLearning-Lecture01.pdf"),
        PyPDFLoader("docs/cs229_lectures/MachineLearning-Lecture02.pdf"),
        PyPDFLoader("docs/cs229_lectures/MachineLearning-Lecture03.pdf")
    ]
    docs = []
    for loader in loaders:
        docs.extend(loader.load())

    # 2、分割文本
    from langchain.text_splitter import RecursiveCharacterTextSplitter
    text_splitter = RecursiveCharacterTextSplitter(
        chunk_size=1500,  # 每个文本块的大小。这意味着每次切分文本时，会尽量使每个块包含 1500 个字符。
        chunk_overlap=150  # 每个文本块之间的重叠部分。
    )
    splits = text_splitter.split_documents(docs)

    # 3、初始化embeddings
    embedding = OpenAIEmbeddings()

    # 4、初始化向量数据库
    persist_directory = 'docs/chroma/cs229_lectures/'
    vectordb = Chroma.from_documents(
        documents=splits,
        embedding=embedding,
        persist_directory=persist_directory  # 允许我们将persist_directory目录保存到磁盘上
    )
    print(vectordb._collection.count()) # 将文档内容拆分为209个块

    # 5、相似性搜索(提问)
    question = "is there an email i can ask for help"  # "有我可以寻求帮助的电子邮件吗"
    docs = vectordb.similarity_search(question, k=3) # k代表检索几个文档
    print("======================")
    print(docs[0].page_content)  # 回答
    print("======================")

    # 6、持久化数据至向量数据库
    vectordb.persist()


if __name__ == '__main__':
    # load_pdf()
    # embeddings_example()
    vectorstores_chroma_example()
